import argparse
import json

"""
为kqa-pro的json文件添加id字段
"""


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", "-i", required=True, help="Input file path")
    parser.add_argument("--output_file", "-o", default=None, help="Output file path")
    args = parser.parse_args()
    if args.output_file is None:
        args.output_file = args.input_file

    with open(args.input_file, "r") as f:
        data = json.load(f)

    file_name = args.input_file.split("/")[-1].split(".")[0]
    if file_name in ["train", "val", "test"]:
        for i, item in enumerate(data):
            item["id"] = f"{file_name}-{i}"

    with open(args.output_file, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4, ensure_ascii=False)


if __name__ == "__main__":
    main()
